Flow Matching For Generative Modeling

This is a learning note from this series of videos.

Link to the paper: https://arxiv.org/abs/2210.02747

Supplementary papers:

Step-by-Step Diffusion: An Elementary Tutorial

Glow: Generative Flow with Invertible 1x1 Convolutions

1. What is Flow? Understand vector field and ODE.

Generative Modeling: With known samples $x_1 \sim q_1 \text{(unknown distribution)}$ , we want to estimate this unknown distribution.

Method: $\underbrace{q_0}_{\text{a known distribution}} \xrightarrow{\phi} \underbrace{q_1}_{\text{unknown distribution to be estimated}}$ . This mapping is denoted as $\phi$ .

How to solve for $\phi$ : 1. Normalizing Flow; 2. Flow Matching (ODE)

Keypoints of Flow Matching:

Let $q_0$ and $q_1$ be the initial and final points of ODE.

Use neural networks to fit the gradient term in ODE.

Solve the ODE.

Preliminaries

ODE → Flow → Normalizing Flow → Continuous Normalizing Flow

Flow

Definition 1. (Flow): A flow is a collection of time-indexed vector fields $v = \{ v_t \}_{t \in [0, 1]}$ .

Any flow defines a trajectory taking initial points $x_1$ to final points $x_0$ , by transporting the initial point along the velocity fields $\{ v_t \}$ . It is equivalent to a transfer between two distributions.

Formally, for velocity field $v$ and initial point $x_1$ , consider the ODE

\begin{equation} \frac{d x_t}{dt} = -v_t(x_t) \end{equation}

with initial condition $x_1$ at time $t=1$ . We write

x_t \coloneqq \text{RunFlow}(v, x_1, t)

to denote the solution to the flow ODE at time $t$ , terminating at final point $x_0$ . That is, RunFlow is the result of transporting point $x_1$ along the flow $v$ up to time $t$ .

Flows also define transports between entire distributions by pushing forward points from the source distribution along their trajectories. If $p_1$ is a distribution on initial points, then applying the flow $v$ yields the distribution on final points:

p_0 = \{ \text{RunFlow}(v, x_1, t=0) \}_{x_1 \sim p_1}

This process is denoted as $p_1 \xrightarrow{v} p_0$ , meaning the flow $v$ transports initial distribution $p_1$ to final distribution $p_0$ .

THE ULTIMATE GOAL OF FLOW MATCHING is to learn a velocity field $v^*$ which transport $p_1 \xrightarrow{v^*} p_0$ , where $p_0$ is the target distribution and $p_1$ is some easy-to-sample base distribution (such as Gaussian).

Continuous Normalizing Flows

Let $\mathbb{R}^d$ denote the data space with data points $x = (x^1, \dots, x^d) \in \mathbb{R}^d$ .

The probability density path $p: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}_{>0}$ is a time-dependent probability density function, i.e., $\int p_t(x) dx = 1$ .

$v: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ is a time-dependent vector field.

A vector field $v_t$ can be used to construct a time-dependent diffeomorphic map, called a flow, $\phi: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ , defined via the ordinary differential equation (ODE):

\begin{align} \frac{d}{dt} \phi_t(x) &= v_t (\phi_t(x)) \\ \phi_0(x) &= x \end{align}

Here $\phi_t(x)$ is a solution to the ODE, and we call it a flow. We can model the vector field $v_t$ with a neural network $v_t(x ; \theta)$ , where $\theta \in \mathbb{R}^p$ are its learnable parameters.

2. The Continuity Equation and Fokker-Planck Equation

Continuity Equation

How to test if a vector field $v_t$ generates a probability path $p_t$ ?

📖

C´edric Villani. Optimal transport: old and new, volume 338. Springer, 2009.

The following PDE provides a necessary and sufficient condition to ensure that a vector field $v_t$ generates a probability path $p_t$ .

\begin{equation} \frac{d}{dt} p_t(x) + \text{div} (p_t(x)v_t(x)) = 0 \end{equation}

Here the divergence operator $\text{div}$ is defined with respect to the spatial variable $x = (x^1, \dots, x^d)$ , i.e. $\text{div} = \sum_{i=1}^d \frac{\partial}{\partial x^i}$ .

Conditional VFs for Fokker-Planck probability paths

Consider a Stochastic Differential Equation (SDE) of the standard form:

\begin{equation} dy = f_t dt + g_t dw \end{equation}

with time parameter $t$ , drift $f_t$ , diffusion coefficient $g_t$ , and $dw$ is the Wiener process.

The solution to the SDE is $y_t$ , which is a stochastic process (a continuous time-dependent variable). Its probability density $p_t(y_t)$ is characterized by the Fokker-Planck equation:

\begin{equation} \frac{dp_t}{dt} = -\text{div}(f_tp_t) + \frac{g_t^2}{2} \Delta p_t \end{equation}

where $\Delta$ represents the Laplace operator (in $y$ ), namely $\text{div} \nabla$ , where $\nabla$ is the gradient operator. We can rewrite the above equation in the form of the continuity equation:

\begin{equation} \begin{align*} \frac{dp_t}{dt} &= -\text{div}(f_t p_t) + \text{div}(\frac{g_t^2}{2}\nabla p_t) \\ &= -\text{div} (f_t p_t - \frac{g_t^2}{2}\frac{\nabla p_t}{p_t} p_t) \\ &= -\text{div}((f_t - \frac{g_t^2}{2} \nabla \log p_t)p_t) \\ &= -\text{div}(w_t p_t) \end{align*} \end{equation}

where the vector field

w_t = f_t - \frac{g_t^2}{2} \nabla \log p_t

satisfies the continuity equation with the probability path $p_t$ , and therefore generates $p_t$ .

3. Continuous Normalizing Flow (CNF)

A CNF is used to reshape a simple prior density $p_0$ (e.g., pure noise) to a more complicated one, $p_1$ , via the push-forward equation.

\begin{equation} p_t = [\phi_t]_* p_0 \end{equation}

The push-forward (or change of variable) operator $*$ is defined by:

\begin{equation} [\phi_t]_* p_0(x) = p_0(\phi_t^{-1}(x)) \det \left[ \frac{\partial \phi_t^{-1}}{\partial x}(x) \right] \end{equation}

📖

Change of variables in the probability density function

Suppose $\bold{x}$ is an n-dimensional random variable with joint density $f$ . If $\bold{y} = G(\bold{x})$ , where $G$ is a bijective, differentiable function, then $\bold{y}$ has density $p_{\bold{Y}}$ :

p_{\bold{Y}}(\bold{y}) = f\left( G^{-1}(\bold{y}) \right) \left| \det \left[ \frac{d G^{-1}(\bold{y})}{d \bold{y}} \right] \right|

Relationships between Flow, Vector Field, and Probability Density

Equation (2)(3) describe the relationship between flow $\phi_t(x)$ and vector field $v_t(x)$

Equation (8)(9) describe how the flow $\phi_t(x)$ changes the density from $p_0$ to $p_t$ .

Equation (4) (continuity equation) gives a necessary and sufficient condition to test whether a vector field $v_t$ generates a probability path $p_t$ .